#include "Task.h"
#include "Utils.h"
#include "Const.h"
#include "Table.h"
#include "Debug.h"
#include "List.h"
#include "Irq.h"
#include "Sem.h"
#include "Mutex.h"
#include "Vm.h"
#include "Syscall.h"
#include "LibSyscall.h"
#include "Mmu.h"
#include "Fs.h"
#include "utility/Instruction.h"
#include "utility/File.h"

static TaskManagerType taskManager;
static uint8_t idleTaskStack[IDLE_TASK_STACK_SIZE];
static TaskType taskTable[TASK_MAX_NUM];
static MutexType taskTableMutex;

void HandleIdleTask(void)
{
    for (;;)
    {
        Hlt();
    }
}

uint32_t InitTask(TaskType *task, const char *name, uint32_t eip, uint32_t esp, bool isKernelTask)
{
    ASSERT(task != (TaskType *)0);

    uint32_t kernelStack = AllocateKernelPage();
    if (kernelStack==0)
    {
        return -1;
    }

    int tssSelector = AllocGdtDesc();
    if (tssSelector < 0)
    {
        Log("GDT is full.");
    }

    SetGDTSegmentDesc(tssSelector, (uint32_t)&task->tss, sizeof(task->tss), SEG_DPL_0 | SEG_P | SEG_TYPE_TSS);

    // Determine whether it is a kernel task
    int codeSelector, dataSelector;
    if (isKernelTask)
    {
        codeSelector = GDT_CODE_SELECTOR;
        dataSelector = GDT_DATA_SELECTOR;
    }
    else
    {
        codeSelector = taskManager.appCodeSelector | SEG_RPL_3;
        dataSelector = taskManager.appDataSelector | SEG_RPL_3;
    }

    Memset(&task->tss, 0, sizeof(task->tss));

    task->processId = (uint32_t)task;

    task->tss.eip = eip;
    task->tss.esp = esp ? esp : kernelStack + PAGE_SIZE;
    task->tss.esp0 = kernelStack + PAGE_SIZE;
    task->tss.ss0 = GDT_DATA_SELECTOR;
    task->tss.ss = task->tss.ds = task->tss.es = task->tss.fs = task->tss.gs = dataSelector;
    task->tss.cs = codeSelector;
    task->tss.eflags = EFLAGS_DEFAULT | EFLAGS_IF;
    task->tss.iomap = 0;

    uint32_t pde = CreateTaskPageDirectoryTable();
    if (pde==0)
    {
        FreeGDTSelector(tssSelector);
        return -1;
    }

    task->tss.cr3 = pde;
    task->tssSelector = tssSelector;

    // A char array in an expression will automatically decay to a char* pointing to the start of the array.
    Strncpy(task->taskName, name, TASK_NAME_SIZE);
    task->state = CREATED;
    InitNode(&task->readyNode);
    InitNode(&task->allNode);
    InitNode(&task->waitNode);

    task->ticks = TASK_TICKS;
    task->remainTicks = TASK_TICKS;
    task->sleepTicks = 0;

    uint32_t eflags = DisableInterrupts();
    InsertNodeAtEnd(&taskManager.taskList, &task->allNode);
    //SetTaskReady(task);
    EnableInterrupts(eflags);
    return 0;
}

// MARK: BUG FIXED
void UninitTask()
{
    TaskType* task = GetCurrentTask();
    if (task->tssSelector)
    {
        FreeGDTSelector(task->tssSelector);
    }
    
    if (task->tss.esp0)
    {
        FreeMemoryPage(task->tss.esp0 - PAGE_SIZE);
    }
    
    if (task->tss.cr3)
    {
        /* code */
    }
    
    Memset(task, 0, sizeof(TaskType));
}

void SwitchTask(TaskType *from, TaskType *to)
{
    FarJump(to->tssSelector, 0);
}

void InitTaskManager()
{
    Memset(taskTable, 0, sizeof(taskTable));
    InitMutex(&taskTableMutex);

    int sel = AllocGdtDesc();
    // Setting the descriptor privilege level (DPL) of a GDT entry
    SetGDTSegmentDesc(sel, 0x00000000, 0xFFFFFFFF, SEG_DPL_3 | SEG_P | SEG_S_NORMAL | SEG_TYPE_DATA | SEG_TYPE_READ_WRITE | SEG_D_32);
    taskManager.appDataSelector = sel;

    sel = AllocGdtDesc();
    SetGDTSegmentDesc(sel, 0x00000000, 0xFFFFFFFF, SEG_DPL_3 | SEG_P | SEG_S_NORMAL | SEG_TYPE_CODE | SEG_TYPE_READ_WRITE | SEG_D_32);
    taskManager.appCodeSelector = sel;

    InitList(&taskManager.taskList);
    InitList(&taskManager.readyTaskList);
    InitList(&taskManager.sleepTaskList);
    InitTask(&taskManager.idleTask, "IdleTask", (uint32_t)HandleIdleTask, 0, true);
    taskManager.currentTask = (TaskType*)0;
    StartTask(&taskManager.idleTask);
}

void InitFirstTask(void)
{
    void FirstTaskEntry(void);
    // Get virtual address
    uint32_t firstTaskEIP = (uint32_t)FirstTaskEntry;

    extern uint8_t firstTaskStart[], firstTaskEnd[];
    uint32_t size = 10 * PAGE_SIZE;
    uint32_t copySize = (uint32_t)(firstTaskEnd - firstTaskStart);

    // eip, esp
    InitTask(&taskManager.firstTask, "FirstTask", firstTaskEIP, firstTaskEIP+size, false);
    taskManager.currentTask = &taskManager.firstTask;
    SetPageDirectoryTable(taskManager.firstTask.tss.cr3);
    AllocateMemoryForCurrentTask(firstTaskEIP, size, PTE_P | PTE_W | PTE_U);
    // After enable paging, virtual address will transfer to physical address automatically
    Memcpy((void*)firstTaskEIP, (void*)firstTaskStart, copySize);

    StartTask(&taskManager.firstTask);
    WriteTR(taskManager.firstTask.tssSelector);
}

TaskType *GetFirstTask()
{
    return &taskManager.firstTask;
}

/// @brief The GetFirstReadyTask function returns a pointer to the first ready task from the readyTaskList managed by taskManager. 
/// If the list is empty, it returns a pointer to the idle task; otherwise, it retrieves the first node from the list and determines its parent task of type TaskType.
TaskType *GetFirstReadyTask()
{
    if (GetListCount(&taskManager.readyTaskList) == 0)
    {
        return &taskManager.idleTask;
    }

    ListNodeType *node = GetFirstNode(&taskManager.readyTaskList);
    return NodeParent(node, TaskType, readyNode);
}

TaskType *GetCurrentTask()
{
    return taskManager.currentTask;
}

static TaskType* AllocateTask(void)
{
    TaskType *task = (TaskType*)0;
    LockMutex(&taskTableMutex);
    for (int i = 0; i < TASK_NAME_SIZE; i++)
    {
        TaskType* curr = &taskTable[i];
        if (curr->taskName[0] == 0)
        {
            task = curr;
            break;
        }
    }
    UnlockMutex(&taskTableMutex);
    return task;
}

void StartTask(TaskType* task)
{
    uint32_t eflags = DisableInterrupts();
    SetTaskReady(task);
    EnableInterrupts(eflags);
}

static void FreeTask(TaskType *task)
{
    LockMutex(&taskTableMutex);
    task->taskName[0] = 0;
    UnlockMutex(&taskTableMutex);
}

/// @brief Set task int ready list
/// @param task
void SetTaskReady(TaskType *task)
{
    if (task != &taskManager.idleTask)
    {
        task->state = READY;
        InsertNodeAtEnd(&taskManager.readyTaskList, &task->readyNode);
    }
}

/// @brief Remove task from ready list
/// @param task
void SetTaskBlock(TaskType *task)
{
    if (task != &taskManager.idleTask)
    {
        task->state = WAITING;
        RemoveNode(&taskManager.readyTaskList, &task->readyNode);
    }
}

/// @brief Insert the task into the sleep queue
/// @param task 
/// @param ticks 
void SetTaskSleep(TaskType *task, uint32_t ticks)
{
    if (ticks > 0)
    {
        task->sleepTicks = ticks;
        task->state = SLEEPING;
        InsertNodeAtEnd(&taskManager.sleepTaskList, &task->readyNode);
    }
}

void SetTaskWakeUp(TaskType *task)
{
    RemoveNode(&taskManager.sleepTaskList, &task->readyNode);
}

void SetTaskWait(TaskType *task, SemaphoreType *sem)
{
    task->state = WAITING;
    InsertNodeAtEnd(&sem->waitList, &task->waitNode);
}

/// @brief Get current task and next task, if they are not equal, set next task as current task
/// @param
void DispatchTask(void)
{
    TaskType *currentTask = GetCurrentTask();
    TaskType *next = GetFirstReadyTask();
    if (taskManager.currentTask != next)
    {
        taskManager.currentTask = next;
        next->state = RUNNING;
        SwitchTask(currentTask, next);
    }
}

/// @brief TaskTimeTick is a function that manages task scheduling in a system by decrementing the remaining ticks for the current task and updating its state when necessary.
///  It also processes sleeping tasks, waking them up when their sleep duration expires, and ensures proper task dispatching while handling interrupt states.

void TaskTimeTick(void)
{
    uint32_t eflags = DisableInterrupts();
    TaskType *task = GetCurrentTask();
    if (--task->remainTicks <= 0)
    {
        task->remainTicks = task->ticks;
        SetTaskBlock(task);
        SetTaskReady(task);
    }

    ListNodeType *sleepNode = GetFirstNode(&taskManager.sleepTaskList);
    while (sleepNode)
    {
        TaskType *sleepTask = NodeParent(sleepNode, TaskType, readyNode);
        if (--sleepTask->sleepTicks <= 0)
        {
            SetTaskWakeUp(sleepTask);
            SetTaskReady(sleepTask);
        }
        sleepNode = sleepNode->next;
    }

    DispatchTask();
    EnableInterrupts(eflags);
}

void SysYield(void)
{
    uint32_t eflags = DisableInterrupts();
    if (GetListCount(&taskManager.readyTaskList) > 1)
    {
        TaskType *task = GetCurrentTask();
        SetTaskBlock(task);
        SetTaskReady(task);
        DispatchTask();
    }
    EnableInterrupts(eflags);
}

void SysSleep(uint32_t ms)
{
    if (ms < TICK_MS)
    {
        ms = TICK_MS;
    }

    uint32_t eflags = DisableInterrupts();

    SetTaskBlock(taskManager.currentTask);
    SetTaskSleep(taskManager.currentTask, (ms + TICK_MS - 1) / TICK_MS);
    DispatchTask();

    EnableInterrupts(eflags);
}

int SysGetPid()
{
    return GetCurrentTask()->processId;
}

int SysFork(void)
{
    TaskType* parentTask = GetCurrentTask();
    TaskType* childTask = AllocateTask();
    if (childTask == (TaskType*)0)
    {
        goto ForkFailed;
    }
    
    // Get system call frame
    // 系统调用发生时，内核会保存用户进程的寄存器现场到内核栈
    // 当用户进程通过系统调用（如调用门或中断）进入内核时，内核会把所有寄存器、参数等信息按固定顺序压入内核栈，形成一个“现场帧”（frame）。
    // esp0 指向内核栈顶
    // parent_task->tss.esp0 是该进程的内核栈顶（即当前内核栈的最高地址）。
    // 现场帧（syscall_frame_t）是刚刚压入栈的内容，位于栈顶之下。
    // frame 的位置
    // 所以，frame 就在 esp0 - sizeof(syscall_frame_t) 这个位置。
    // 通过强制类型转换，可以直接用指针访问这段内存，把它当作 syscall_frame_t 结构体来操作。
    SyscallFrameType* frame = (SyscallFrameType*)(parentTask->tss.esp0 - sizeof(SyscallFrameType));

    // Init fork task
    // 在调用Syscall的时候，会封存当前的寄存器情况，其中esp指向的是最后传入的参数，所以这里要加上参数个数*4字节。
    
    int error = InitTask(childTask, parentTask->taskName, frame->eip, frame->esp + sizeof(uint32_t) * SYSCALL_PARAM_COUNT, false);
    if (error<0)
    {
        goto ForkFailed;
    }

    TssDescType* tss = &childTask->tss;
    tss->eax = 0;                       
    tss->ebx = frame->ebx;
    tss->ecx = frame->ecx;
    tss->edx = frame->edx;
    tss->esi = frame->esi;
    tss->edi = frame->edi;
    tss->ebp = frame->ebp;

    tss->cs = frame->cs;
    tss->ds = frame->ds;
    tss->es = frame->es;
    tss->fs = frame->fs;
    tss->gs = frame->gs;
    tss->eflags = frame->eflags;

    childTask->parent = parentTask;
    childTask->tss.cr3 = CopyUserVirtualMemoryAndPageTable(parentTask->tss.cr3);
    if (childTask->tss.cr3<0)
    {
        goto ForkFailed;
    }

    StartTask(childTask);
    return childTask->processId;

ForkFailed:
    if (childTask)
    {
        FreeTask(childTask);
    }
    return -1;
}

static int LoadPhdr(int file, Elf32_Phdr* phdr, uint32_t pageDir)
{
    int error = AllocateMemory(pageDir, phdr->p_vaddr, phdr->p_memsz, PTE_P | PTE_W | PTE_U);
    if (error)
    {
        return error;
    }

    // cp = targetAddr + p_offset
    if (SyslSeek(file, phdr->p_offset, 0) < 0)
    {
        return -1;
    }
    
    uint32_t vaddr = phdr->p_vaddr;
    uint32_t size = phdr->p_filesz;
    while (size > 0)
    {
        uint32_t count = size > PAGE_SIZE ? PAGE_SIZE : size;
        uint32_t paddr = GetPhysicalAddress(pageDir, vaddr);

        // Read data from cp to paddr
        if (SysRead(file, (char*)paddr, count) < count)
        {
            return -1;
        }
        vaddr += count;
        size -= count;
    }
    return 0;
}


static uint32_t LoadElfFlile(TaskType* task, const char* name, uint32_t pageDir)
{
    Elf32_Ehdr ehdr;
    Elf32_Phdr phdr;

    // cp = targetAddress
    int file = SysOpen(name, 0);
    if (file<0)
    {
        goto LoadFailed;
    }
    
    // Read ELF header data
    // cp = targetAddr + sizeof(ehdr)
    int count = SysRead(file, (char*)&ehdr, sizeof(Elf32_Ehdr));
    if (count < sizeof(Elf32_Ehdr))
    {
        goto LoadFailed;
    }
    
    if (ehdr.e_ident[0] != ELF_IDENTIFY || ehdr.e_ident[1] != 'E' || 
        ehdr.e_ident[2] != 'L' || ehdr.e_ident[3] != 'F')
    {
        goto LoadFailed;
    }
    
    if (ehdr.e_type != ET_EXEC || ehdr.e_machine != ET_386 || ehdr.e_entry==0)
    {
        goto LoadFailed;
    }
    
    if (ehdr.e_phentsize==0 || ehdr.e_phoff==0)
    {
        goto LoadFailed;
    }

    uint32_t e_phoff = ehdr.e_phoff;
    for (int i = 0; i < ehdr.e_phnum; i++, e_phoff += ehdr.e_phentsize)
    {
        // Set file pointer to the current phdr
        // cp = targetAddr + e_phoff + e_phentsize * i.  phdr table address
        if (SyslSeek(file, e_phoff, 0) < 0)
        {
            goto LoadFailed;
        }

        // Load data to phdr
        // cp += sizeof(phdr)
        count = SysRead(file, (char*)&phdr, sizeof(Elf32_Phdr));
        if (count < sizeof(Elf32_Phdr))
        {
            goto LoadFailed;
        }
        
        if (phdr.p_type != PT_LOAD || phdr.p_vaddr < TASK_ADDRESS_BASE)
        {
            continue;
        }

        int error = LoadPhdr(file, &phdr, pageDir);
        if (error < 0)
        {
            goto LoadFailed;
        }
    }

    return ehdr.e_entry;

LoadFailed:
    if (file>=0)
    {
        SysClose(file);
    }

    return 0;
}


static int CopyArgs(char* to, uint32_t pageDir, int argc, char** argv)
{
    TaskArgsType taskArgs;
    taskArgs.argc = argc;
    taskArgs.argv = (char**)(to + sizeof(TaskArgsType));

    // Address of char* array content
    char* charContDest = to + sizeof(TaskArgsType) + argc * sizeof(char*); 
    char** charTableDest = (char**)GetPhysicalAddress(pageDir, (uint32_t)(to + sizeof(TaskArgsType)));

    for (int i = 0; i < argc; i++)
    {
        char* from = argv[i];
        int len = Strlen(from) + 1;
        int error = CopyDataFromOtherPageDir((uint32_t)charContDest, pageDir, (uint32_t)from, len);

        charTableDest[i] = charContDest;
        charContDest += len;
    }
    
    return CopyDataFromOtherPageDir((uint32_t)to, pageDir, (uint32_t)&taskArgs, sizeof(TaskArgsType));
}


int SysExecve(char* name, char* argv[], char* env[])
{
    TaskType* task = GetCurrentTask();

    Strncpy(task->taskName, GetFileName(name), 32);

    uint32_t oldPageDir = task->tss.cr3;
    uint32_t newPageDir = CopyUserVirtualMemoryAndPageTable(oldPageDir);

    if (!newPageDir)
    {
        goto ExecFailed;
    }
    
    uint32_t entry = LoadElfFlile(task, name, newPageDir);
    if (entry==0)
    {
        goto ExecFailed;
    }

    uint32_t stackTop = MEM_TASK_STACK_TOP - MEM_TASK_ARG_SIZE;
    int error = AllocateMemory(newPageDir, 
        MEM_TASK_STACK_TOP - MEM_TASK_STACK_SIZE,
        MEM_TASK_STACK_SIZE, PTE_W | PTE_U | PTE_P);
    if (error < 0)
    {
        goto ExecFailed;
    }

    int argc = StringCount(argv);
    error = CopyArgs((char*)stackTop, newPageDir, argc, argv);
    if (error<0)
    {
        goto ExecFailed;
    }
    
    SyscallFrameType* frame = (SyscallFrameType*)(task->tss.esp0 - sizeof(SyscallFrameType));
    frame->eip = entry;
    frame->eax = frame->ebx = frame->ecx = frame->edx = 0;
    frame->esi = frame->edi = frame->ebp = 0;
    frame->eflags = EFLAGS_DEFAULT | EFLAGS_IF;

    frame->esp = stackTop - sizeof(uint32_t)* SYSCALL_PARAM_COUNT;

    task->tss.cr3 = newPageDir;
    SetPageDirectoryTable(newPageDir);

    FreeUserVirtualMemoryAndPageTable(oldPageDir);
    
    return 0;

ExecFailed:
    if (newPageDir)
    {
        task->tss.cr3 = oldPageDir;
        SetPageDirectoryTable(oldPageDir);
        FreeUserVirtualMemoryAndPageTable(newPageDir);
    }

    return -1;
}


